import network
import socket
import json
import time
from machine import Pin, PWM

# --- Wi-Fi Configuration ---
WIFI_SSID = 'Servo'
WIFI_PASSWORD = '12345678' # Change this! Min 8 characters for WPA2-PSK

# --- Server Configuration ---
SERVER_HOST = '0.0.0.0' # Listen on all available interfaces
SERVER_PORT = 8888

# --- Servo Configuration ---
# !!! IMPORTANT: CALIBRATE THESE VALUES FOR YOUR SERVOS !!!
# These are typical for SG90 servos with 50Hz PWM.
# Duty cycle is for 16-bit resolution (0-65535)
# 0.5ms pulse for 0 degrees: (0.5 / 20) * 65535 = 1638
# 2.5ms pulse for 180 degrees: (2.5 / 20) * 65535 = 8191
SERVO_MIN_DUTY_U16 = 1600  # Duty for 0 degrees
SERVO_MAX_DUTY_U16 = 8200  # Duty for 180 degrees
PWM_FREQUENCY = 50

# --- Servo Pins (CHANGE THESE TO MATCH YOUR WIRING) ---
SERVO_PIN_VERTICAL = 4
SERVO_PIN_HORIZONTAL = 12
SERVO_PIN_FLIP = 14

# --- Global Variables ---
pwm_vertical = None
pwm_horizontal = None
pwm_flip = None

last_angles = {"vertical": 90.0, "horizontal": 90.0, "flip": 90.0}
SMOOTHING_FACTOR = 0.3 # For simple exponential smoothing (0 = no smoothing, 1 = instant)
ANGLE_CHANGE_THRESHOLD = 0.5 # Only update servo if angle changes by more than this (degrees)

# --- Functions ---
def setup_wifi_ap():
    ap = network.WLAN(network.AP_IF)
    ap.active(False) # Deactivate first to apply new config
    time.sleep(0.5)
    ap.config(essid=WIFI_SSID, password=WIFI_PASSWORD, authmode=network.AUTH_WPA2_PSK)
    ap.active(True)
    
    while not ap.active():
        time.sleep(0.1)
        print("Waiting for AP to activate...")
        
    print('ESP32 Access Point Configured')
    print('SSID:', WIFI_SSID)
    print('IP Address:', ap.ifconfig()[0])
    return ap

def init_servos():
    global pwm_vertical, pwm_horizontal, pwm_flip
    
    pin_v = Pin(SERVO_PIN_VERTICAL, Pin.OUT)
    pin_h = Pin(SERVO_PIN_HORIZONTAL, Pin.OUT)
    pin_f = Pin(SERVO_PIN_FLIP, Pin.OUT)
    
    pwm_vertical = PWM(pin_v, freq=PWM_FREQUENCY)
    pwm_horizontal = PWM(pin_h, freq=PWM_FREQUENCY)
    pwm_flip = PWM(pin_f, freq=PWM_FREQUENCY)
    
    # Set servos to initial center position
    set_servo_angle("vertical", 90.0)
    set_servo_angle("horizontal", 90.0)
    set_servo_angle("flip", 90.0)
    print("Servos initialized to 90 degrees.")

def angle_to_duty_u16(angle):
    """Converts an angle (0-180) to a 16-bit PWM duty cycle value."""
    clamped_angle = max(0, min(180, angle))
    return int(SERVO_MIN_DUTY_U16 + (clamped_angle / 180.0) * (SERVO_MAX_DUTY_U16 - SERVO_MIN_DUTY_U16))

def set_servo_angle(servo_name, target_angle):
    global last_angles
    
    # Apply simple exponential smoothing
    current_angle = last_angles.get(servo_name, target_angle)
    smoothed_angle = (target_angle * SMOOTHING_FACTOR) + (current_angle * (1.0 - SMOOTHING_FACTOR))
    
    # Only update if change is significant, to reduce jitter
    if abs(smoothed_angle - current_angle) < ANGLE_CHANGE_THRESHOLD and servo_name in last_angles:
        # print(f"Skipping {servo_name} update, change too small.")
        return

    last_angles[servo_name] = smoothed_angle
    duty = angle_to_duty_u16(smoothed_angle)
    
    if servo_name == "vertical" and pwm_vertical:
        pwm_vertical.duty_u16(duty)
    elif servo_name == "horizontal" and pwm_horizontal:
        pwm_horizontal.duty_u16(duty)
    elif servo_name == "flip" and pwm_flip:
        pwm_flip.duty_u16(duty)
    # print(f"Servo {servo_name} set to {smoothed_angle:.1f}° (duty: {duty})")


def handle_client_connection(client_socket):
    global last_angles
    print("Client connected")
    buffer = ""
    try:
        while True:
            request_bytes = client_socket.recv(1024)
            if not request_bytes:
                break # Connection closed by client
            
            buffer += request_bytes.decode('utf-8')
            
            # Process all newline-terminated messages in the buffer
            while '\n' in buffer:
                message, buffer = buffer.split('\n', 1)
                if not message: # Skip empty messages that might result from multiple newlines
                    continue
                try:
                    # print(f"Received raw: '{message}'")
                    angles_data = json.loads(message)
                    
                    v_angle = float(angles_data.get("vertical", last_angles["vertical"]))
                    h_angle = float(angles_data.get("horizontal", last_angles["horizontal"]))
                    f_angle = float(angles_data.get("flip", last_angles["flip"]))
                    
                    # Clamp angles on ESP32 side as well, just in case
                    v_angle = max(0, min(180, v_angle))
                    h_angle = max(0, min(180, h_angle))
                    f_angle = max(0, min(180, f_angle))

                    set_servo_angle("vertical", v_angle)
                    set_servo_angle("horizontal", h_angle)
                    set_servo_angle("flip", f_angle)
                    
                    print(f"Processed angles: V={v_angle:.1f}, H={h_angle:.1f}, F={f_angle:.1f}")

                except json.JSONDecodeError:
                    print(f"JSON Decode Error for: '{message}'")
                except ValueError:
                    print(f"ValueError (non-float data?) for: '{message}'")
                except Exception as e:
                    print(f"Error processing message: '{message}', Error: {e}")
                    
    except OSError as e:
        print(f"Socket error with client: {e}")
    except Exception as e:
        print(f"Unexpected error with client: {e}")
    finally:
        client_socket.close()
        print("Client disconnected")

def start_server():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) # Allow reusing address
    s.bind((SERVER_HOST, SERVER_PORT))
    s.listen(1) # Listen for one connection at a time
    print(f'TCP Server listening on {SERVER_HOST}:{SERVER_PORT}')

    while True:
        try:
            client_sock, addr = s.accept()
            print(f'Accepted connection from {addr}')
            handle_client_connection(client_sock)
        except KeyboardInterrupt:
            print("Server shutting down by KeyboardInterrupt.")
            break
        except Exception as e:
            print(f"Error in server accept loop: {e}")
            time.sleep(1) # Wait a bit before retrying
    
    if pwm_vertical: pwm_vertical.deinit()
    if pwm_horizontal: pwm_horizontal.deinit()
    if pwm_flip: pwm_flip.deinit()
    s.close()
    print("Server shut down.")

# --- Main Execution ---
if __name__ == "__main__":
    try:
        init_servos()
        ap_lan = setup_wifi_ap()
        start_server()
    except Exception as e:
        print(f"Fatal error in main: {e}")
        # Optional: try to clean up or reset
        if 'ap_lan' in locals() and ap_lan.active():
            ap_lan.active(False)
        if pwm_vertical: pwm_vertical.deinit()
        if pwm_horizontal: pwm_horizontal.deinit()
        if pwm_flip: pwm_flip.deinit()
        print("System halted due to error.")